package com.example.demo.tree;

/**
    声明：本代码非本人所写，全部copy为TreeMap

 */

import java.util.Comparator;
import java.util.HashMap;

public class RBTree<K,V> {

    private final static boolean BLACK = true;
    private final static boolean RED = false;

    private transient Entry<K,V> root;

    private final Comparator<? super K> comparator;

    private transient int size;
    private transient int modCount;
    
    static final class Entry<K,V>{
        public boolean color = BLACK;
        public K key ;
        public V value;
        public Entry<K,V> left;
        public Entry<K,V> right;
        public Entry<K,V> parent;

        public Entry(K key, V value) {
            this.key = key ;
            this.value = value;
        }

        public Entry(K key, V value, Entry<K,V> p) {
            this.key = key ;
            this.value = value;
            this.parent = p;
        }

        public String toString() {
            return "value=" + value.toString();
        }
        
        public K getKey() {
            return this.key;
        }


        public V getValue() {
            return this.value;
        }


        public V setValue(V value) {
            V old = this.value;
            this.value = value;
            return old;
        }
    }

    public RBTree() {
        this.comparator = null;
    }


    public V add(K key,V value) {

        if (root == null) {
            root = new Entry<>(key,value);
            return null;
        }

        Entry<K,V> p = root;
        Entry<K,V> parent;
        int cmp;

        if (comparator != null) {
            Comparator<? super K> cpr = this.comparator;
            do {
                parent = p;
                cmp = cpr.compare(key, parent.key);
                if (cmp == 0) {
                    return parent.setValue(value);
                } else if (cmp < 0) {
                    p = parent.left;
                } else {
                    p = parent.right;
                }
            } while (p != null);
        } else {
            if (value == null) {
                throw new NullPointerException();
            }
            Comparable<? super K> k = (Comparable<? super K>) key;
            do {
                parent = p;
                cmp = k.compareTo(parent.key);
                if (cmp == 0) {
                    return parent.setValue(value);
                } else if (cmp < 0) {
                    p = parent.left;
                } else {
                    p = parent.right;
                }
            } while (p != null);
        }


        Entry<K,V> newNode = new Entry<>(key,value, parent);

        if (cmp < 0) {
            parent.left = newNode;
        } else {
            parent.right = newNode;
        }

        fixAfterInsertion(newNode);
        size++;
        modCount++;
        return null;
    }


    private void fixAfterInsertion(Entry<K,V> x) {
        x.color = RED;

        while (x != null && x != root && x.parent.color == RED) {
            // 父亲节点是祖父节点的左孩子 ​​​​
            if (parentOf(x) == leftOf(parentOf(parentOf(x)))) {
                //  叔父节点
                Entry<K,V> y = rightOf(parentOf(parentOf(x)));
                //  叔父节点是红色
                if (colorOf(y) == RED) {
                    // 父节点设置为黑色
                    setColor(parentOf(x), BLACK);
                    // 叔父节点设置为黑色
                    setColor(y, BLACK);
                    // 祖父节点设置为黑色
                    setColor(parentOf(parentOf(x)), RED);
                    //
                    x = parentOf(parentOf(x));
                } else {
                    //  叔父节点是黑色
                    //   x节点是父亲节点的右孩子  LR型，先调整为LL型
                    if (x == rightOf(parentOf(x))) {
                        x = parentOf(x);
                        rotateLeft(x);
                    }
                    //   x节点是父亲节点的左孩子
                    //   父节点设置为黑色
                    setColor(parentOf(x), BLACK);
                    //  祖父节点设置为红色
                    setColor(parentOf(parentOf(x)), RED);
                    //  插入过程只需要进行一次旋转
                    rotateRight(parentOf(parentOf(x)));
                }
            } else {
                Entry<K,V> y = leftOf(parentOf(parentOf(x)));
                if (colorOf(y) == RED) {
                    setColor(parentOf(x), BLACK);
                    setColor(y, BLACK);
                    setColor(parentOf(parentOf(x)), RED);
                    x = parentOf(parentOf(x));
                } else {
                    if (x == leftOf(parentOf(x))) {
                        x = parentOf(x);
                        rotateRight(x);
                    }
                    setColor(parentOf(x), BLACK);
                    setColor(parentOf(parentOf(x)), RED);
                    rotateLeft(parentOf(parentOf(x)));
                }
            }
        }
        root.color = BLACK;
    }


    private static <K,V> boolean colorOf(Entry<K,V> p) {
        return (p == null ? BLACK : p.color);
    }

    private static <K,V> Entry<K,V> parentOf(Entry<K,V> p) {
        return (p == null ? null: p.parent);
    }

    private static <K,V> void setColor(Entry<K,V> p, boolean c) {
        if (p != null)
            p.color = c;
    }

    private static <K,V> Entry<K,V> leftOf(Entry<K,V> p) {
        return (p == null) ? null: p.left;
    }

    private static <K,V> Entry<K,V> rightOf(Entry<K,V> p) {
        return (p == null) ? null: p.right;
    }

    private void rotateLeft(Entry<K,V> p) {
        if (p != null) {
            Entry<K,V> r = p.right;
            p.right = r.left;
            if (r.left != null)
                r.left.parent = p;
            r.parent = p.parent;
            if (p.parent == null)
                root = r;
            else if (p.parent.left == p)
                p.parent.left = r;
            else
                p.parent.right = r;
            r.left = p;
            p.parent = r;
        }
    }

    /** From CLR */
    private void rotateRight(Entry<K,V> p) {
        if (p != null) {
            Entry<K,V> l = p.left;
            p.left = l.right;
            if (l.right != null) l.right.parent = p;
            l.parent = p.parent;
            if (p.parent == null)
                root = l;
            else if (p.parent.right == p)
                p.parent.right = l;
            else p.parent.left = l;
            l.right = p;
            p.parent = l;
        }
    }

    public V remove(Object key) {
        Entry<K,V> p = getEntry(key);
        if (p == null)
            return null;

        V oldValue = p.value;
        deleteEntry(p);
        return oldValue;
    }

    final Entry<K,V> getEntry(Object key) {
        // Offload comparator-based version for sake of performance
        if (comparator != null)
            return getEntryUsingComparator(key);
        if (key == null)
            throw new NullPointerException();
        @SuppressWarnings("unchecked")
        Comparable<? super K> k = (Comparable<? super K>) key;
        Entry<K,V> p = root;
        while (p != null) {
            int cmp = k.compareTo(p.key);
            if (cmp < 0)
                p = p.left;
            else if (cmp > 0)
                p = p.right;
            else
                return p;
        }
        return null;
    }


    final Entry<K,V> getEntryUsingComparator(Object key) {
        @SuppressWarnings("unchecked")
        K k = (K) key;
        Comparator<? super K> cpr = comparator;
        if (cpr != null) {
            Entry<K,V> p = root;
            while (p != null) {
                int cmp = cpr.compare(k, p.key);
                if (cmp < 0)
                    p = p.left;
                else if (cmp > 0)
                    p = p.right;
                else
                    return p;
            }
        }
        return null;
    }


    private void deleteEntry(Entry<K,V> p) {
        modCount++;
        size--;

        // If strictly internal, copy successor's element to p and then make p
        // point to successor.
        if (p.left != null && p.right != null) {
            Entry<K,V> s = successor(p);
            p.key = s.key;
            p.value = s.value;
            p = s;
        } // p has 2 children

        // Start fixup at replacement node, if it exists.
        Entry<K,V> replacement = (p.left != null ? p.left : p.right);

        if (replacement != null) {
            // Link replacement to parent
            replacement.parent = p.parent;
            if (p.parent == null)
                root = replacement;
            else if (p == p.parent.left)
                p.parent.left  = replacement;
            else
                p.parent.right = replacement;

            // Null out links so they are OK to use by fixAfterDeletion.
            p.left = p.right = p.parent = null;

            // Fix replacement
            if (p.color == BLACK)
                fixAfterDeletion(replacement);
        } else if (p.parent == null) { // return if we are the only node.
            root = null;
        } else { //  No children. Use self as phantom replacement and unlink.
            if (p.color == BLACK)
                fixAfterDeletion(p);

            if (p.parent != null) {
                if (p == p.parent.left)
                    p.parent.left = null;
                else if (p == p.parent.right)
                    p.parent.right = null;
                p.parent = null;
            }
        }
    }
    
    // 寻找后续节点
    static <K,V> Entry<K,V> successor(Entry<K,V> t) {
        if (t == null)
            return null;
        else if (t.right != null) {
            Entry<K,V> p = t.right;
            while (p.left != null)
                p = p.left;
            return p;
        } else {
            Entry<K,V> p = t.parent;
            Entry<K,V> ch = t;
            while (p != null && ch == p.right) {
                ch = p;
                p = p.parent;
            }
            return p;
        }
    }


    private void fixAfterDeletion(Entry<K,V> x) {

        // x 节点的颜色为黑色
        while (x != root && colorOf(x) == BLACK) {
            // x 节点为左孩子
            if (x == leftOf(parentOf(x))) {
                // sib为右孩子
                Entry<K,V> sib = rightOf(parentOf(x));
                // 兄弟节点为红色
                if (colorOf(sib) == RED) {
                    setColor(sib, BLACK);
                    setColor(parentOf(x), RED);
                    // 向左旋转。 像黑色节点那边旋转
                    rotateLeft(parentOf(x));
                    //  sib 指向兄弟节点
                    sib = rightOf(parentOf(x));
                }

                // sib 为黑色节点
                //  sib的左右节点都为黑色
                if (colorOf(leftOf(sib))  == BLACK &&
                        colorOf(rightOf(sib)) == BLACK) {
                    //  sib 设置为红色
                    setColor(sib, RED);
                    //  回溯x指向x的父亲节点
                    x = parentOf(x);
                } else {
                    //  右孩子节点颜色为黑色
                    if (colorOf(rightOf(sib)) == BLACK) {
                        // 左节点也设置为黑色
                        setColor(leftOf(sib), BLACK);
                        // sib 设置为红色
                        setColor(sib, RED);
                        //  向右旋转
                        rotateRight(sib);
                        //  sib 指向右孩子
                        sib = rightOf(parentOf(x));
                    }
                    //
                    setColor(sib, colorOf(parentOf(x)));
                    setColor(parentOf(x), BLACK);
                    setColor(rightOf(sib), BLACK);
                    rotateLeft(parentOf(x));
                    x = root;
                }
            } else { // symmetric
                Entry<K,V> sib = leftOf(parentOf(x));

                if (colorOf(sib) == RED) {
                    setColor(sib, BLACK);
                    setColor(parentOf(x), RED);
                    rotateRight(parentOf(x));
                    sib = leftOf(parentOf(x));
                }

                if (colorOf(rightOf(sib)) == BLACK &&
                        colorOf(leftOf(sib)) == BLACK) {
                    setColor(sib, RED);
                    x = parentOf(x);
                } else {
                    if (colorOf(leftOf(sib)) == BLACK) {
                        setColor(rightOf(sib), BLACK);
                        setColor(sib, RED);
                        rotateLeft(sib);
                        sib = leftOf(parentOf(x));
                    }
                    setColor(sib, colorOf(parentOf(x)));
                    setColor(parentOf(x), BLACK);
                    setColor(leftOf(sib), BLACK);
                    rotateRight(parentOf(x));
                    x = root;
                }
            }
        }

        setColor(x, BLACK);
    }

    public boolean checkInvariants() {
        return checkInvariants(root);
    }

    public static <K,V> boolean checkInvariants(Entry<K,V> t) {

        if (t == null) {
            return true;
        }

        if (t.parent != null || t.color == RED) {
            return false;
        }

        int lc = check(t, t.left);
        int rc = check(t, t.right);

        if(lc==rc){
            if(lc == 0){
                return false ;
            }
            System.out.println(lc);
            return true ;

        }
        return false ;
    }


    static <K,V> int check(Entry<K,V> p,Entry<K,V> t){

        // 假设叶子节点为黑色
        if (t == null) {
            return 1;
        }

        if (t.parent != p) {
            return 0;
        }

        // 红不能连接红色
        if(p.color == RED && t.color == RED){
            return 0 ;
        }

        //  t到每个后裔叶子节点(black nil node) 经历相同的黑色节点
        int lc = check(t, t.left);
        int rc = check(t, t.right);

        if (lc == rc) {
            if (lc == 0) {
                return 0;
            }

            if (t.color == BLACK) {
                return lc + 1;
            }
            return lc;
        }
        return 0;
    }


}
